//
//  binarySearchTree.swift
//  SetAndMap
//
//  Created by lhc on 2024/6/22.
//

import Foundation

class BinarySearchTree<E: Comparable> {
    
    enum Exception: Error {
        case emptyTree
        case notContainElement(E)
    }
    
    class Node: Equatable {
        static func == (lhs: BinarySearchTree<E>.Node, rhs: BinarySearchTree<E>.Node) -> Bool {
            return lhs.e == rhs.e
        }
        
        
        var e: E
        var left: Node?
        var right: Node?
        
        init(e: E) {
            self.e = e
        }
    }
    
    private var root: Node? = nil
    private var size = 0
    
    /// 获取二分搜索树size
    func getSize() -> Int {
        return self.size
    }
    
    /// 判断二分搜索树是否为空
    func isEmpty() -> Bool {
        return self.size == 0
    }
    
    ///  二分搜索树添加新元素(递归方式外部调用)
    ///  - 参数: e: Node承载的元素
    func add(e: E) {
        self.root = self.add(self.root, e: e)
    }
    
    /// 二分搜索树内部实际的添加元素方法(递归调用)
    ///  - node: 传入的node
    ///  - e: 添加的元素
    ///
    private func add(_ node: Node?, e: E) -> Node {
        guard let node = node else{
            // 如果传入的node为空 用e生成一个新node返回给父节点
            self.size += 1
            return Node(e: e)
        }
        
        if e < node.e {
            node.left = self.add(node.left, e: e)
        }else if e > node.e {
            node.right = self.add(node.right, e: e)
        }
        
        // 如果node.e == e
        // 因为二分搜索树没有重复元素,所以什么都不做
            
        return node
    }
    
    /// 二分搜索树添加新元素(循环的方式)
    func cycleAdd(e: E) {
        guard let root = self.root else{
            // 如果根节点为空, 直接生成新的根节点
            self.size += 1
            self.root = Node(e: e)
            return
        }
        
        var n = root
        // 停止条件传入的e值二分搜索树里已经有了
        while n.e != e {
            if e < n.e {
                if n.left == nil {
                    // 当n的左子树为空时,新创建node元素安排在左子树
                    n.left = Node(e: e)
                    self.size += 1
                    break
                }else {
                    // 左子树不为空, n指向n的左子树继续循环
                    n = n.left!
                }
            }else if e > n.e {
                if n.right == nil {
                    // 当n的右子树为空时,新创建node元素安排在右子树
                    n.right = Node(e: e)
                    self.size += 1
                    break
                }else {
                    // 右子树不为空, n指向n的右子树继续循环
                    n = n.right!
                }
            }
        }
    }
    
    /// 查找是否包含元素(外部调用)
    func containsElement(_ element: E) -> Bool {
        return self.containsElement(element, node: self.root)
    }
    
    
    /// 查找是否包含元素(递归调用)
    private func containsElement(_ element: E, node: Node?) -> Bool {
        guard let node = node else {
            return false
        }
        
        if element < node.e {
            return self.containsElement(element, node: node.left)
        }else if element > node.e {
            return self.containsElement(element, node: node.right)
        }else {
            return true
        }
    }
    
    /// 前序遍历(递归调用, 外部调用)
    func preOrder() {
        self.preOrder(self.root)
    }
    
    /// 前序遍历(递归调用)
    private func preOrder(_ node: Node?) {
        guard let node = node else { return }
        
        print(node.e)
        self.preOrder(node.left)
        self.preOrder(node.right)
    }
    
    /// 中序遍历(递归调用, 外部调用)
    func inOrder() {
        self.inOrder(self.root)
    }
    
    /// 中序遍历(递归调用)
    private func inOrder(_ node: Node?) {
        guard let node = node else { return }
        
        self.inOrder(node.left)
        print(node.e)
        self.inOrder(node.right)
    }
    
    /// 后序遍历(递归调用, 外部调用)
    func postOrder() {
        self.postOrder(self.root)
    }
    
    /// 后序遍历(递归调用)
    private func postOrder(_ node: Node?) {
        guard let node = node else { return }
        
        self.postOrder(node.left)
        self.postOrder(node.right)
        print(node.e)
    }
    
    /// 前序遍历(非递归调用, 外部调用)
    func preOrderNoRecursion() {
        // 利用链表栈实现前序遍历
        let linklistStack = LinkListStack<Node>()
        guard !self.isEmpty() else { return }
        
        // 因为上面判断了二分搜索树不为空, 所以传入的root一定不是nil
        linklistStack.push(e: self.root!)
        
        // 循环判断的终止条件是链表栈不为空
        while !linklistStack.isEmpty() {
            // 取出栈顶node打印
            let node = linklistStack.pop()
            print(node?.e ?? "nil")
            // 如果有右孩子放入栈顶
            if let right = node?.right {
                linklistStack.push(e: right)
            }
            // 如果有左孩子放入栈顶
            if let left = node?.left {
                linklistStack.push(e: left)
            }
            // 先右后左入栈, 取的时候就是先左后右
        }
    }
    
    /// 层序遍历
    func levelOrder(action: (E?) -> ()) {
        // 利用链表队列实现层序遍历
        let linklistQueue = LinkListQueue<Node>()
        
        guard !self.isEmpty() else { return }
        
        linklistQueue.enqueue(e: self.root!)
        
        while !linklistQueue.isEmpty() {
            let node = linklistQueue.dequeue()
            
            action(node?.e)
            
            if let left = node?.left {
                linklistQueue.enqueue(e: left)
            }
            
            if let right = node?.right {
                linklistQueue.enqueue(e: right)
            }
        }
        
    }
    
    /// 返回最小值的结点(外部调用)
    func miniNode() -> Node? {
        guard !self.isEmpty() else { return nil }
        
        return self.miniNode(self.root!)
    }
    
    /// 返回最小值的结点
    private func miniNode(_ node: Node) -> Node? {
        if let left = node.left {
            return self.miniNode(left)
        }else {
            return node
        }
    }
    
    /// 返回最大值的结点(外部调用)
    func maxNode() -> Node? {
        guard !self.isEmpty() else { return nil }
        
        return self.maxNode(self.root!)
    }
    
    /// 返回最大值的结点
    private func maxNode(_ node: Node) -> Node? {
        if let right = node.right {
            return self.maxNode(right)
        }else {
            return node
        }
    }
    
    /// 删除最小值结点(外部调用)
    func removeMiniNode() -> Node? {
        guard !self.isEmpty() else { return nil }
        // 先找到最小node保存，后面删除后返回
        let node = self.miniNode()
        
        self.root = self.removeMiniNode(self.root!)
        
        return node
    }
    
    /// 删除最小值结点
    private func removeMiniNode(_ node: Node) -> Node? {
        guard let left = node.left else {
            // 找到要删除的最小值就是当前node
            
            // 保存右子结点，不管是不是空
            let rightChildNode = node.right
            // 释放node持有的右子结点, 防止内存泄漏
            node.right = nil
            self.size -= 1
            
            // 把先前保存的右子结点返回回去替换父结点的左子结点
            // 这样父节点就没有持有应该删除的最小值node， node就被删掉了
            return rightChildNode
        }
        
        node.left = self.removeMiniNode(left)
        return node
    }
    
    /// 删除最大值结点(外部调用)
    func removeMaxNode() -> Node? {
        guard !self.isEmpty() else { return nil }
        // 先找到最大值node保存，后面在树中删除后返回
        let node = self.maxNode()
        
        self.root = self.removeMaxNode(self.root!)
        
        return node
    }
    
    /// 删除最大值结点
    private func removeMaxNode(_ node: Node) -> Node? {
        guard let right = node.right else {
            // 找到要删除的node了, 因为他没有右子树了
            
            // 先保存要删除的node的左子树，不管他是不是空
            // 后面要返回给父节点的作为右子树
            let leftChildNode = node.left
            
            // node的左结点指针设置为空
            node.left = nil
            
            self.size -= 1
            
            // 把先前保存的左子结点返回回去替换父结点的右子结点
            // 这样父节点就没有持有应该删除的最大值node， node就被删掉了
            return leftChildNode
        }
        
        node.right = self.removeMaxNode(right)
        return node
    }
    
    /// 移除元素
    func removeElement(_ element: E) throws {
        guard !self.isEmpty() else {
            throw Exception.emptyTree
        }
        
        self.root = try self.removeElement(self.root, element: element)
    }
    
    private func removeElement(_ node: Node?, element: E) throws -> Node?  {
        // 如果node为空, 那么返回一个空给父节点, 毕竟父节点的子结点为空返回一个空给他重新接上没问题
        guard let node = node else {
            throw Exception.notContainElement(element)
        }
        
        if element < node.e {
            // 目标元素比当前node的元素小，继续在node的左子树寻找
            node.left = try self.removeElement(node.left, element: element)
            return node
        }else if element > node.e {
            // 目标元素比当前node的元素大，继续在node的右子树寻找
            node.right = try self.removeElement(node.right, element: element)
            return node
        }
        
        // 剩下的情况就是找到了element == node.e
        
        // 如果没有左子树
        if node.left == nil {
            // 保存右子树作为返回node，可以为空
            let rightChildNode = node.right
            // 删除node的右结点指针
            node.right = nil
            
            self.size -= 1
            return rightChildNode
        }
        
        // 如果没有右子树
        if node.right == nil {
            let leftChildNode = node.left
            node.left = nil
            self.size -= 1
            return leftChildNode
        }
        
        // 剩下的情况就是又有左子树又有右子树
        // 选择当前结点的大一点的继结点(就是右子树上的最小结点)
        let successor = self.miniNode(node.right!)
        // 继结点的右结点绑定上当前结点的右结点
        successor?.right = self.removeMiniNode(node.right!)
        successor?.left = node.left
        
        node.left = nil
        node.right = nil
        // 把successor返回给父节点替代当前node node就被删除了, 因为在removeMiniNode方法里有size -= 1所以不用专门写size -= 1
        return successor
    }
}
    
